package org.xiaoyu.rpc.registry.zookeeper;

import cn.hutool.core.collection.CollectionUtil;
import lombok.extern.slf4j.Slf4j;
import org.apache.curator.RetryPolicy;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.framework.recipes.cache.PathChildrenCache;
import org.apache.curator.framework.recipes.cache.PathChildrenCacheListener;
import org.apache.curator.retry.ExponentialBackoffRetry;
import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.data.Stat;
import org.xiaoyu.rpc.common.constant.SystemPropertiesKey;
import org.xiaoyu.rpc.common.exception.RegistryException;
import org.xiaoyu.rpc.common.url.URL;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

/**
 * Created By Have
 * 2021/7/15 9:50
 */
@Slf4j
public class ZookeeperInstanceWrapper {
    private final static String DEFAULT_ROOT = "/rpc";
    private final static int DEFAULT_MAX_RETRY = 3;
    private final static int DEFAULT_SLEEP_TIME = 1000;

    private final String root;

    // todo add node listener
    private static final Map<String, List<String>> SERVICE_MAP = new ConcurrentHashMap<>();

    private final CuratorFramework zkClient;
    private final int maxRetry;
    private final int sleepTime;
    private final String serverAddr;

    public ZookeeperInstanceWrapper(Properties properties) {
        if (properties == null) {
            throw new RegistryException("Registry properties cannot be null");
        }
        this.root = properties.getProperty(SystemPropertiesKey.ZK_ROOT) != null ? properties.getProperty(SystemPropertiesKey.ZK_ROOT): DEFAULT_ROOT;
        this.maxRetry = properties.getProperty(SystemPropertiesKey.REGISTRY_MAX_RETRY) != null ?  Integer.valueOf(properties.getProperty(SystemPropertiesKey.REGISTRY_MAX_RETRY)) : DEFAULT_MAX_RETRY;
        this.sleepTime = properties.getProperty(SystemPropertiesKey.SLEEP_TIME) != null? Integer.valueOf(properties.getProperty(SystemPropertiesKey.SLEEP_TIME)) : DEFAULT_SLEEP_TIME;
        this.serverAddr = properties.getProperty(SystemPropertiesKey.REGISTER_ADDR);
        if (serverAddr == null || serverAddr.isEmpty()) {
            throw new RegistryException("Registry address can not be empty or null");
        }

        RetryPolicy retryPolicy = new ExponentialBackoffRetry(this.sleepTime, this.maxRetry);
        this.zkClient = CuratorFrameworkFactory.builder()
                .connectString(serverAddr)
                .retryPolicy(retryPolicy)
                .build();
        zkClient.start();
        try {
            if (!zkClient.blockUntilConnected(30, TimeUnit.SECONDS)) {
                throw new RegistryException(String.format("Timeout to connect ZK server with address [%s]", serverAddr));
            }
        } catch (InterruptedException e) {
            if (log.isErrorEnabled()) {
                log.error("Cannot connect ZK Server, Reason: [{}]", e.getMessage());
                throw new RegistryException(e);
            }
        }

        if (log.isInfoEnabled()) {
            log.info("Connect zk Server success with address[{}]", serverAddr);
        }
    }

    public void registerService(URL url) {
        createPersistentNode(url);
    }

    public void createPersistentNode(URL url) {
        String path = this.root + url.getAllPath();
        try {
            Stat stat = zkClient.checkExists().forPath(path);
            if (stat != null) {
                return;
            }
            zkClient.create().creatingParentsIfNeeded().withMode(CreateMode.EPHEMERAL).forPath(path);
            log.info("The node with path [{}] has been registered", path);
        } catch (Exception e) {
            log.info("Register node failed, Reason [{}]", e.getMessage());
            throw new RegistryException(e);
        }
    }

    public void unRegistered(URL url) {
        try {
            zkClient.delete().forPath(this.root + url.getAllPath());
        } catch (Exception e) {
            log.error("Registry delete path failed, Reason:[{}]", e.getMessage());
            throw new RegistryException(e);
        }
    }


    public List<String> lookup(URL url) {
        String servicePath = this.root + url.getServicePath();
        List<String> result = SERVICE_MAP.get(servicePath);
        if (CollectionUtil.isNotEmpty(result)) {
            return result;
        }
        try {
            registerWatcher(url);
            result = zkClient.getChildren().forPath(servicePath);
        } catch (Exception e) {
            log.error("zookeeper search child node fail: [{}]", e.getMessage());
        }
        return result;
    }

    public List<String> lookupStrictly(URL url) {
        String strictPath = this.root + url.getAllPath();
        List<String> result = new ArrayList<>();
        try {
            result = zkClient.getChildren().forPath(strictPath);
        } catch (Exception e) {
            log.error("zookeeper strictly search child node with path [{}] fail, Reason: [{}]", strictPath, e.getMessage());
        }
        return result;
    }

    public void disconnect() {
        this.zkClient.close();
    }

    private void registerWatcher(URL url) throws Exception {
        String servicePath = this.root + url.getServicePath();
        PathChildrenCache pathChildrenCache = new PathChildrenCache(zkClient, servicePath, true);
        PathChildrenCacheListener pathChildrenCacheListener = (curatorFramework, pathChildrenCacheEvent) -> {
            List<String> serviceAddress = curatorFramework.getChildren().forPath(servicePath);
            SERVICE_MAP.put(servicePath, serviceAddress);
        };
        pathChildrenCache.getListenable().addListener(pathChildrenCacheListener);
        pathChildrenCache.start();
    }
}
